# -*- coding: utf-8 -*-
"""
Created on Fri Dec 18 12:10:02 2020

@author: 26297
"""

#  模型预加载数据的预处理部分
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader
from age_net import AgeNet
import matplotlib.pyplot as plt
import torch.nn.functional as F
import time
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 处理器选择
print(device)
# device = torch.device('cuda')
print("模型加载...")
path = '../Arcface_100.pth'
agenet = AgeNet.agenet(path)
PATH = '../model/MSE_model_001.pth'
agenet.load_state_dict(torch.load(PATH))

print("数据预处理...")
# 数据预处理
transform = transforms.Compose([
    transforms.Resize(112),
    transforms.ToTensor()
])

# 读取数据
root = '../age_image/train1/img_bust2'
train_dataset = datasets.ImageFolder(root, transform)
# test_dataset = datasets.ImageFolder(root + '/test', transform)

# 导入数据
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=8, shuffle=True,num_workers=8)
# test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=8, shuffle=True)

classes = train_dataset.classes
classes_index = train_dataset.class_to_idx

# 参数全部参与训练
count = 472
for param in agenet.parameters():
    count -= 1
    param.requires_grad = True
    # # 将Arcface中的参数改为不可变
    # if(count>4):
    #     param.requires_grad = False
    # # 将自定义模型中的参数改为可变
    # else:
    #     param.requires_grad = True

#   模型训练和测试部分

# 参数设置
LR = 0.01  # 学习率

# 第二类代价函数，平方差损失函数，计算二范数
MSE_loss = nn.MSELoss().to(device)
# 定义优化器
optimizer = optim.Adam(agenet.parameters(), LR)


class Trainer:
    __init__(self):
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 处理器选择
        # device = torch.device('cuda')
        
        print("模型加载...")
        path = '../Arcface_100.pth'
        agenet = AgeNet.agenet(path)
        PATH = '../model/MSE_model_001.pth'
        agenet.load_state_dict(torch.load(PATH))





def train(train_loader, start_epoch, end_epoch, batch):
    # 读取数据
    test_root = '../age_image/test/img_bust2'
    test_dataset = datasets.ImageFolder(root, transform)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=8, shuffle=True)

    agenet.train()

    loss_list = []  # 记录损失值的变化
    acc_list = []  # 记录准去率的变化
    bias = 0
    for i, data in enumerate(train_loader):
        print("第{}个batch".format(i), end='\t')
        t0 = time.time()
        # 获得数据和对应的标签
        inputs, labels = data
        inputs.to(device)
        inputs.requires_grad = True
        
        t1 = time.time()
        # 获得模型预测结果
        out = agenet(inputs).view(labels.size())
        # 计算交叉熵LOss
        # loss= entropy_loss(out, labels)
        t2 = time.time()
        # 计算二范数loss
        loss = MSE_loss(out.float(), labels.float()).to(device)
        loss_list.append(loss.item())
        t3 = time.time()
        print("loss: ", loss.item())
        print("real:{}  pre:{}".format(labels, out))
        # print("real:{}  pre:{}".format(labels, torch.max(out, 1).indices))
        
        # 计算梯度
        loss.backward()
        t4 = time.time()
        
        bias += (out-labels).sum()
        
        # 修改权值
        optimizer.step()
        # 每训练50次计算一次准确率

        # 梯度清0
        optimizer.zero_grad()
        t5 = time.time()
        # 每个epoch的训练次数为batch
        if i >= batch:
            break
        print('{0:.4f}-{1:.4f}-{2:.4f}-{3:.4f}-{4:.4f}'.format((t1-t0),(t2-t1), (t3-t2),(t4-t3),(t5-t4)))

    torch.save(agenet.state_dict(), '../model/MSE_model_001.pth')
    return loss_list,bias/i/8
        # 将损失值、准确率记录到文件中，用于绘制变化趋势图
        # with open("./loss_record.txt", "a+") as f:
        #     f.write("第{}个epoch".format(ep))
        #     for loss in loss_list:
        #         f.write(str(loss) + '\n')

        # with open("./acc_record.txt", "a+") as f:
        #     f.write("第{}个epoch".format(ep))
        #     for acc in acc_list:
        #         f.write(str(acc) + '\n')
        


def MSE_acc(test_loader, left, right):
    correct = 0
    total = 0
    bias = 0
    for i, data in enumerate(test_loader):
        agenet.eval()
        if i < 15:
            inputs, labels = data
            inputs.requires_grad = False
            out = agenet(inputs).view(labels.size())
            bias += (out - labels).sum()
            # out = out.round()
            # correct += (out == labels).sum()
            out = (left <= out - labels) * (out - labels <= right)

            correct += out.sum()

            total += len(labels)


        else:
            break
    acc = correct.item() / total
    return acc, bias / total


if __name__ == '__main__':
    start_epoch = 25
    end_epoch = 30
    batch = 5000  # 每个轮次的训练次数
    loss_record, bias = train(train_loader, start_epoch, end_epoch, batch)
    x1 = [i for i in range(len(loss_record))]
    plt.plot(x1, loss_record)
    plt.show()
#
# # 读取数据
# test_root = '../age_image/test/img_bust2'
# test_dataset = datasets.ImageFolder(root, transform)
# test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=8, shuffle=True)
#
# for ep in range(start_epoch, epoch):
#     print('第{}个epoch:'.format(ep))
#     loss_list = []  # 记录损失值的变化
#     acc_list = []  # 记录准去率的变化
#     train(train_loader, test_loader, loss_list, acc_list)
#
#     torch.save(agenet.state_dict(), '../model/MSE_model_{}.pth'.format(ep))
#     # 将损失值、准确率记录到文件中，用于绘制变化趋势图
#     with open("./loss_record.txt", "a+") as f:
#         for loss in loss_list:
#             f.write(str(loss) + '\n')
#
#     with open("./acc_record.txt", "a+") as f:
#         for acc in acc_list:
#             f.write(str(acc) + '\n')
#
#     if ep%2==0:
#         # 展示损失值变化图
#         loss_record = []
#         for loss in open("./loss_record.txt"):
#             loss_record.append(float(loss))
#
#         acc_record = []
#         for acc in open("./acc_record.txt"):
#             acc_record.append(float(acc))
#
#         # 展示损失值、准确率的变化
#         x1 = [i for i in range(len(loss_record))]
#         plt.plot(x1, loss_record)
#         x2 = [i for i in range(len(acc_record))]
#         plt.plot(x2, acc_record)
#         plt.show()
